{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# nanoAhaMoment: Single File \"RL for LLM\" Library\n",
    "Single GPU · No TRL or Verl · Efficient · 3B Base Model · Full Parameter Tuning Implementation of R1-zero training.\n",
    "\n",
    "Inspired by [TinyZero](https://github.com/Jiayi-Pan/TinyZero) and [Mini-R1](https://www.philschmid.de/mini-deepseek-r1), but designed to be **simpler**, **cleaner**, and **faster**, with every line of code visible and understandable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "R1-Zero is arguably the more interesting contribution from the DeepSeek R1 paper. The core idea: take a freshly pre-trained LLM (straight out of the unsupervised pretraining oven) and continue its training using reinforcement learning *without* any human feedback or supervision. The result? A model that starts showing emergent behaviors like self-reflection, verification, backtracking that researchers have tried to bake into LLMs using handcrafted tricks and inductive biases, at least since O1.\n",
    "\n",
    "In this notebook, we’ll build an R1-Zero-style training loop **from scratch**. The goal is to create a crystal-clear, hackable foundation for RL-style LLM training; one that gives you a bird’s-eye view of every moving part and how they fit together. Perfect for playing around, extending, or hacking.\n",
    "\n",
    "---\n",
    "\n",
    "### Why another R1-Zero implementation?\n",
    "\n",
    "There are already great implementations like [TinyZero](https://github.com/Jiayi-Pan/TinyZero) and [Mini-R1](https://www.philschmid.de/mini-deepseek-r1). But they rely on full-fledged RL libraries (like `trl` or `verl`) to handle training.\n",
    "\n",
    "These libraries exist for good reason; efficient RL training for LLMs sits at the crossroads of scalable training and fast inference. Making that work takes a lot of engineering. But that also means the internals are often abstracted away, hard to read, and even harder to tweak.\n",
    "\n",
    "This notebook is different: **no abstractions, no hiding**. You’ll see everything, top to bottom. A lightweight, readable codebase that still follows best practices and runs efficiently on a single GPU.\n",
    "\n",
    "### What is this notebook, exactly?\n",
    "\n",
    "We'll train a base LLM using RL to solve a reasoning-heavy algorithmic task. The setup:\n",
    "\n",
    "- **Model**: Qwen2.5 3B-Base  \n",
    "- **Dataset**: Countdown-Tasks-3to4  \n",
    "- **Algorithm**: GRPO (a variant of policy gradient)\n",
    "\n",
    "Yes, the task is a bit toy-ish—but it captures the essence of R1-Zero: emergent behaviors like self-reflection, verification, backtracking, even language-switching. This setup is ideal for rapid prototyping and experimentation.\n",
    "\n",
    "### Who is this notebook for?\n",
    "\n",
    "- Anyone interested in RL training for LLMs  \n",
    "- Researchers, especially the ones in academia, exploring reasoning in language models\n",
    "\n",
    "### What should I know before jumping in?\n",
    "\n",
    "- A working knowledge of the HuggingFace Transformers library  \n",
    "- Some experience fine-tuning LLMs  \n",
    "- Familiarity with policy gradient methods (helpful but not required)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## R1-Zero Recipe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal is to train a base LLM to **reason** in a way that allows it to **reevaluate** its own outputs and **improve** them, all without human supervision. The DeepSeek R1 paper proposes a surprisingly simple recipe to achieve this, and that's exactly what we'll implement in this notebook.\n",
    "\n",
    "### The Recipe\n",
    "\n",
    "Here's the high-level procedure:\n",
    "\n",
    "1. **Start** with a base LLM and a dataset containing problem prompts paired only with their *final answers* (no intermediate reasoning steps).  \n",
    "2. For each iteration $i = 0$ to `NUM_ITERATIONS`:\n",
    "   - Sample a batch of prompts $\\{x_i\\}_{i=1}^N$ from the dataset.\n",
    "   - For each prompt, sample $G$ responses from the model:  \n",
    "     $ y_1, y_2, \\cdots, y_G \\sim \\pi_\\theta(y|x) $\n",
    "\n",
    "     These $G$ responses form what is called a *group* in GRPO.\n",
    "   - Compute a reward $R_i$ for each response and normalize them tocalculate the GRPO advantage within each group.\n",
    "   - Create a list of $N \\times G$ episodes, i.e., pairs of $(x_i, y_i)$ along with their corresponding advantages.\n",
    "   - Estimate the policy gradient $\\vec{g}_{pg}$ from these episodes.\n",
    "   - Update the model parameters:  \n",
    "     $\\theta \\leftarrow \\theta + \\eta \\vec{g}_{pg}$\n",
    "\n",
    "### Code Structure Overview\n",
    "\n",
    "The code you will see is structured directly following this recipe. It boils down to three main components:\n",
    "\n",
    "1. **Episode Generation**  \n",
    "   - Generate $ (x, y) $ pairs along with their advantages for each RL iteration.\n",
    "   \n",
    "2. **Reward Calculation**  \n",
    "   - Compute rewards for each generated response.\n",
    "   \n",
    "3. **Policy Gradient Estimation**  \n",
    "   - Use the generated episodes to estimate the policy gradient and perform the model update.\n",
    "\n",
    "In the end, these three components come together in a simple loop that trains the model, step by step, to develop reasoning capabilities through reinforcement learning.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checkpoint Playground"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the `notebooks/checkpoint_playground.ipynb`, you can load the model we already trained with this notebook and interactively test the model's reasoning capabilities. This notebook allows you to input custom prompts and observe the model's responses."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Installing Dependencies\n",
    "\n",
    "Before we begin, let's install the necessary Python packages. We'll be using:\n",
    "\n",
    "- PyTorch  \n",
    "- Hugging Face Transformers  \n",
    "- Hugging Face Datasets  \n",
    "- DeepSpeed  \n",
    "- vLLM\n",
    "\n",
    "For a detailed, step-by-step installation guide, refer to the [README](https://github.com/McGill-NLP/tiny-aha-moment.git) of this project."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "# Set the environment variables for HuggingFace\n",
    "# This is done to ensure that the cache directory for HuggingFace is set to a specific location,\n",
    "# preventing the storage from being overwhelmed with model files and other data.\n",
    "SCRATCH = Path.home() / \"scratch\"\n",
    "os.environ[\"HF_HOME\"] = str(SCRATCH / \"hf_home\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "import re\n",
    "import time\n",
    "from typing import Any, Dict, List, Tuple, Union\n",
    "\n",
    "import deepspeed\n",
    "import numpy as np\n",
    "import torch\n",
    "from datasets import load_dataset\n",
    "from deepspeed import DeepSpeedEngine\n",
    "from tqdm import trange\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel\n",
    "from vllm import LLM, SamplingParams\n",
    "\n",
    "import wandb\n",
    "from utils import (\n",
    "    compute_token_log_probs,\n",
    "    dump_episodes,\n",
    "    evaluate_on_test_set,\n",
    "    find_free_port,\n",
    "    find_last_checkpoint,\n",
    "    prepare_model_inputs,\n",
    "    load_model_into_vllm\n",
    ")\n",
    "\n",
    "# Needed to stop DeepSpeed from complaining\n",
    "os.environ[\"MASTER_ADDR\"] = \"localhost\"\n",
    "os.environ[\"MASTER_PORT\"] = str(find_free_port())\n",
    "os.environ[\"RANK\"] = \"0\"\n",
    "os.environ[\"LOCAL_RANK\"] = \"0\"\n",
    "os.environ[\"WORLD_SIZE\"] = \"1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**We do have a few helper functions in `utils.py` that are used to keep the code clean.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyperparameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's define the hyperparameters for the training. These are mostly taken from [Mini-R1](https://www.philschmid.de/mini-deepseek-r1) implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model configuration\n",
    "MODEL_NAME = \"Qwen/Qwen2.5-3B\"\n",
    "MODEL_CHAT_NAME = MODEL_NAME + \"-Instruct\"\n",
    "\n",
    "# Dataset configuration\n",
    "DATASET_NAME = \"Jiayi-Pan/Countdown-Tasks-3to4\"\n",
    "\n",
    "# Total number of training iterations\n",
    "NUM_ITERATIONS = 1000\n",
    "# Number of episodes to collect per iteration for training\n",
    "EPISODES_PER_ITERATION = 64\n",
    "# Number of responses to generate for each input prompt (i.e. group size in GRPO)\n",
    "GENERATIONS_PER_SAMPLE = 4\n",
    "# Controls how much the policy can deviate from the reference model\n",
    "KL_COEFFICIENT = 0.001\n",
    "\n",
    "# Training hyperparameters\n",
    "# Batch size for each GPU device during training\n",
    "PER_DEVICE_BATCH_SIZE = 4\n",
    "# Learning rate for model updates\n",
    "LEARNING_RATE = 1e-6\n",
    "\n",
    "# Sampling parameters\n",
    "# Maximum number of tokens to generate in each response\n",
    "MAX_RESPONSE_TOKENS = 1024\n",
    "# Controls randomness in generation (higher = more random)\n",
    "TEMPERATURE = 1.0\n",
    "# Nucleus sampling parameter (1.0 = disabled)\n",
    "TOP_P = 1.0\n",
    "# Top-k sampling parameter (-1 = disabled)\n",
    "TOP_K = -1  # no top k\n",
    "\n",
    "# DeepSpeed configuration\n",
    "# DeepSpeed config for the policy model\n",
    "deepspeed_config = {\n",
    "    \"bf16\": {\"enabled\": True},\n",
    "    \"zero_optimization\": {\"stage\": 2, \"overlap_comm\": False},\n",
    "    \"train_batch_size\": EPISODES_PER_ITERATION,\n",
    "    \"train_micro_batch_size_per_gpu\": PER_DEVICE_BATCH_SIZE,\n",
    "    \"gradient_accumulation_steps\": EPISODES_PER_ITERATION // PER_DEVICE_BATCH_SIZE,\n",
    "    \"gradient_clipping\": 1.0,\n",
    "    \"optimizer\": {\n",
    "        \"type\": \"AdamW\",\n",
    "        \"params\": {\n",
    "            \"lr\": LEARNING_RATE,\n",
    "            \"betas\": (0.9, 0.999),\n",
    "            \"eps\": 1e-8,\n",
    "            \"weight_decay\": 0.0,\n",
    "            \"torch_adam\": True,\n",
    "        },\n",
    "    },\n",
    "}\n",
    "# DeepSpeed config for the reference model\n",
    "ref_deepspeed_config = {\n",
    "    \"bf16\": {\"enabled\": True},\n",
    "    # Note that we don't train the reference model\n",
    "    # These are just for compatibility with DeepSpeed.\n",
    "    \"train_batch_size\": EPISODES_PER_ITERATION,\n",
    "    \"train_micro_batch_size_per_gpu\": PER_DEVICE_BATCH_SIZE,\n",
    "    \"gradient_accumulation_steps\": EPISODES_PER_ITERATION // PER_DEVICE_BATCH_SIZE,\n",
    "}\n",
    "\n",
    "RUN_NAME = \"r1-zero\"\n",
    "EXP_DIR = SCRATCH / \"deepseek_r1z_hackathon\" / RUN_NAME\n",
    "EXP_DIR.mkdir(parents=True, exist_ok=True)\n",
    "print(f\"Logs and Checkpoints will be saved to: {EXP_DIR}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating the training prompts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For training, we'll use the [Countdown-Tasks-3to4](https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4) dataset, which provides problem statements paired with their final answers (but no reasoning steps).\n",
    "\n",
    "### The Countdown Task\n",
    "\n",
    "The Countdown game is a numerical puzzle where the player must reach a target number using a set of randomly chosen numbers and basic arithmetic operations: addition, subtraction, multiplication, and division. Each number must be used exactly once.\n",
    "\n",
    "Example:\n",
    "\n",
    "```yaml\n",
    "Target: 622\n",
    "Available Numbers: [25, 3, 6, 100]\n",
    "\n",
    "# Not provided in the dataset\n",
    "Solution: (100 × 6) + (25 − 3) = 622\n",
    "```\n",
    "\n",
    "This task is ideal for training LLMs to practice reasoning, searching, and self-verification.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are using the base version of the model, which has only been pretrained on raw internet data, it has no prior understanding of system prompts or chat formatting. However, we will still use the chat format to make the resulting model compatible with downstream tools and frameworks that expect it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_MESSAGE = (\n",
    "    \"You are a helpful assistant. You first think about the reasoning process in the mind \"\n",
    "    \"and then provide the user with the answer.\"\n",
    ")\n",
    "PROMPT_TEMPLATE = (\n",
    "    \"Using the numbers {numbers}, create an equation that equals {target}. \"\n",
    "    \"You can use basic arithmetic operations (+, -, *, /) and each number can only be used once. \"\n",
    "    \"Show your work in <think> </think> tags. And return the final equation and answer in \"\n",
    "    \"<answer> </answer> tags, for example <answer>(1 + 2) / (3 * 5)</answer>.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the system message and prompt template, we can generate the training prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and process dataset\n",
    "def preprocess_example(example: Dict[str, Any]):\n",
    "    numbers: List[int] = example[\"nums\"]\n",
    "    target: int = example[\"target\"]\n",
    "\n",
    "    prefix = [\n",
    "        {\"role\": \"system\", \"content\": SYSTEM_MESSAGE},\n",
    "        {\"role\": \"user\", \"content\": PROMPT_TEMPLATE.format(numbers=numbers, target=target)},\n",
    "        {\"role\": \"assistant\", \"content\": \"Let me solve this step by step.\\n<think>\"},\n",
    "    ]\n",
    "    input_ids = tokenizer.apply_chat_template(\n",
    "        prefix, tokenize=True, continue_final_message=True\n",
    "    )\n",
    "    prompt = tokenizer.decode(\n",
    "        input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False\n",
    "    )\n",
    "    return {\"prompt\": prompt, \"input_ids\": input_ids}\n",
    "\n",
    "# Note that the base model and \"instruct\" model have different eos token. \n",
    "# Here we make sure to use the correct one.\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_CHAT_NAME)\n",
    "EOS_TOKEN_ID = AutoTokenizer.from_pretrained(MODEL_NAME).eos_token_id\n",
    "EOS_TOKEN = tokenizer.convert_ids_to_tokens(EOS_TOKEN_ID)\n",
    "\n",
    "dataset = load_dataset(DATASET_NAME, split=\"train\")\n",
    "dataset = dataset.map(preprocess_example, num_proc=6)\n",
    "\n",
    "# Split dataset\n",
    "train_test_split = dataset.train_test_split(test_size=500, seed=42)\n",
    "train_dataset = train_test_split[\"train\"]\n",
    "test_dataset = train_test_split[\"test\"]\n",
    "\n",
    "len(train_dataset), len(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's look at some examples from the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Target: \", train_dataset[0][\"target\"])\n",
    "print(\"Available Numbers: \", train_dataset[0][\"nums\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the system message and prompt template, we generate the following prompt for this example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_dataset[0][\"prompt\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you noticed, we also prepend the `<assistant>` tag along with the phrase *\"Let me solve this step by step.\"* to each prompt. This helps guide the model into **answering mode**. Without this, the base model might simply continue the prompt rather than attempting to solve the task, since it has no inherent understanding of instruction-following.\n",
    "\n",
    "Additionally, we tokenize each prompt and store the result as `input_ids`, which will be used later during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train_dataset[0][\"input_ids\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reward Function\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The DeepSeek R1 paper introduced **rule-based rewards** to evaluate whether the model-generated solutions were correct. We'll adopt a similar approach by defining two custom reward functions:\n",
    "\n",
    "- **Format Reward**: Checks if the output follows the required format:  \n",
    "  `<think> [thinking] </think><answer> [answer] </answer>`\n",
    "\n",
    "- **Equation Reward**: Extracts the equation from within the `<answer>` tag, verifies that it evaluates to the target result, and ensures that all available numbers are used exactly once.\n",
    "\n",
    "The purpose of enforcing the format is mainly to make answer extraction easier. It isn't strictly necessary for the correctness of the answer itself but simplifies parsing during training.\n",
    "\n",
    "The final reward assigned to an episode/trajectory (prompt+response) is simply the sum of these two components. Importantly, the reward is only computed at the **last token** of the output. From an RL perspective, this means that all intermediate actions receive zero reward. We also do not apply any discounting here (i.e., $\\gamma = 1$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_reward_func(completion: str) -> float:\n",
    "    \"\"\"\n",
    "    Format: <think>...</think>\\n</answer>...</answer>\n",
    "\n",
    "    Also checks that the content within <answer>...</answer> conforms to a\n",
    "    specified pattern (only digits, + - * / ( ) . and whitespace).\n",
    "\n",
    "    Args:\n",
    "        completion (str): Generated output\n",
    "\n",
    "    Returns:\n",
    "        float: Reward score\n",
    "    \"\"\"\n",
    "    # Define the allowed pattern (only numbers, +, -, *, /, (, ), ., and whitespace)\n",
    "    allowed_pattern = r\"^[\\d+\\-*/().\\s]+$\"\n",
    "\n",
    "    try:\n",
    "        # add synthetic <think> as its already part of the prompt and prefilled \n",
    "        # for the assistant to more easily match the regex\n",
    "        completion = \"<think>\" + completion\n",
    "\n",
    "        # Strip EOS token if present\n",
    "        if completion.endswith(EOS_TOKEN):\n",
    "            completion = completion[:-len(EOS_TOKEN)]\n",
    "\n",
    "        # Check if the format is correct\n",
    "        # Pattern means:\n",
    "        # 1) <think>...contents not including other <think> tags...</think>\n",
    "        # 2) \\n\n",
    "        # 3) <answer>...anything...</answer>\n",
    "        regex = r\"^<think>([^<]*(?:<(?!/?think>)[^<]*)*)<\\/think>\\n<answer>([\\s\\S]*?)<\\/answer>$\"\n",
    "        match = re.search(regex, completion, re.DOTALL)\n",
    "\n",
    "        if match is None or len(match.groups()) != 2:\n",
    "            # Format is incorrect\n",
    "            return 0.0\n",
    "        else:\n",
    "            # Extract the content inside <answer>...</answer>\n",
    "            answer_content = match.group(2).strip()\n",
    "\n",
    "            # Check if answer content matches the allowed pattern\n",
    "            if not re.match(allowed_pattern, answer_content):\n",
    "                # If it doesn't match, reward is 0.5\n",
    "                return 0.5\n",
    "            else:\n",
    "                # If both format and pattern are correct, reward is 1\n",
    "                return 1.0\n",
    "    except Exception:\n",
    "        # Any error leads to 0 reward\n",
    "        return 0.0\n",
    "\n",
    "\n",
    "def equation_reward_func(completion: str, nums: List[int], target: int) -> float:\n",
    "    \"\"\"\n",
    "    Evaluates completion based on mathematical correctness of the answer\n",
    "\n",
    "    Args:\n",
    "        completion (str): Generated output\n",
    "        target (str): Expected answer\n",
    "        nums (list): Available numbers to use in the equation\n",
    "\n",
    "    Returns:\n",
    "        float: Reward score\n",
    "    \"\"\"\n",
    "    try:\n",
    "        # Check if the format is correct\n",
    "        match = re.search(r\"<answer>(.*?)<\\/answer>\", completion)\n",
    "        if match is None:\n",
    "            return 0.0\n",
    "        # Extract the \"answer\" part from the completion\n",
    "        equation = match.group(1).strip()\n",
    "        # Extract all numbers from the equation\n",
    "        used_numbers = [int(n) for n in re.findall(r\"\\d+\", equation)]\n",
    "\n",
    "        # Check if all numbers are used exactly once\n",
    "        if sorted(used_numbers) != sorted(nums):\n",
    "            return 0.0\n",
    "        # Define a regex pattern that only allows numbers, operators, parentheses, and whitespace\n",
    "        allowed_pattern = r\"^[\\d+\\-*/().\\s]+$\"\n",
    "        if not re.match(allowed_pattern, equation):\n",
    "            return 0.0\n",
    "\n",
    "        # Evaluate the equation with restricted globals and locals\n",
    "        result = eval(equation, {\"__builtins__\": None}, {})\n",
    "        # Check if the equation is correct and matches the ground truth\n",
    "        if abs(float(result) - float(target)) < 1e-5:\n",
    "            return 1.0\n",
    "        else:\n",
    "            return 0.0\n",
    "    except Exception:\n",
    "        # If evaluation fails, reward is 0\n",
    "        return 0.0\n",
    "    \n",
    "\n",
    "def compute_reward(completion: str, sample: Dict[str, Any]) -> Tuple[float, Dict[str, float]]:\n",
    "    nums = sample[\"nums\"]\n",
    "    target = sample[\"target\"]\n",
    "\n",
    "    format_reward = format_reward_func(completion)\n",
    "    equation_reward = equation_reward_func(\n",
    "        completion=completion, nums=nums, target=target\n",
    "    )\n",
    "\n",
    "    reward = format_reward + equation_reward\n",
    "\n",
    "    metrics = {\n",
    "        \"format_reward\": format_reward,\n",
    "        \"equation_reward\": equation_reward,\n",
    "    }   \n",
    "\n",
    "    return reward, metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# <think> is prefilled in the prompt. So, repeating it in the completion would be incorret.\n",
    "format_reward_func(\"<think>I think the answer is </think>\\n<answer>1+2</answer>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "format_reward_func(\"I think the answer is </think>\\n<answer>1+2</answer>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "format_reward_func(\"<think>I think the<think>and even more</think> answer is </think>\\n<answer>1+2</answer>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "equation_reward_func(\"I think the answer is </think>\\n<answer>1+2+2</answer>\", [1,2], 3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Episode Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal of episode generation is to create a collection of query-response pairs that will be used for policy training. From the reinforcement learning (RL) perspective, the **query** serves as the initial state, and the generated tokens in the **response** represent the actions taken by the policy.\n",
    "\n",
    "The `create_training_episodes` function takes a list of prompts (initial states) and their corresponding completions which we generate using the model.  In GRPO, we always generate multiple responses per prompt—specifically, `GENERATIONS_PER_SAMPLE` > 1. This means that, after episode generation, we end up with `batch_size × GENERATIONS_PER_SAMPLE` episodes in every RL iteration.\n",
    "\n",
    "### Advantage Computation\n",
    "\n",
    "In addition to generating episodes, `create_training_episodes` is also responsible for computing the **advantage** for every response token. \n",
    "\n",
    "In RL terms, the advantage of a token represents how much better or worse that token's action is compared to the average generate token at that specific state (prompt + prefix). Ideally, we would compute an advantage for every token individually to capture how each step contributes to the overall reward.\n",
    "\n",
    "However, in GRPO, there's no per-token advantage computation. Instead, we compute a single advantage value per response. This value reflects how good the entire response is relative to other responses generated for the same prompt. We then assign this single advantage value uniformly to all tokens within that response.\n",
    "\n",
    "GRPO uses a simple formula for this:\n",
    "\n",
    "1. For each prompt $x$ with a group of generated responses $y_1, y_2, \\ldots, y_G \\sim \\pi(\\cdot|x)$, compute their rewards $R_1, R_2, \\ldots, R_G$.\n",
    "2. Compute the group's mean and standard deviation:  \n",
    "   $ \\mu = \\text{mean}(R_1, R_2, \\ldots, R_G) $  \n",
    "   $ \\sigma = \\text{std}(R_1, R_2, \\ldots, R_G) $\n",
    "3. Compute a **relative score** for each response:  \n",
    "   $ R^*_i = \\frac{R_i - \\mu}{\\sigma} $\n",
    "4. Assign this relative score $R^*_i$ as the advantage to all tokens of the $i$-th response:  \n",
    "   $ A_t^{(i)} = R^*_i $\n",
    "\n",
    "This **per-group normalization** encourages responses that are better than average and penalizes those that are worse.\n",
    "\n",
    "### Example: Advantage in Action\n",
    "\n",
    "Consider a binary reward scenario where each response is either correct (1) or incorrect (0):\n",
    "\n",
    "```python\n",
    ">>> rewards = np.array([1, 1, 0, 0, 0])\n",
    ">>> (rewards - rewards.mean()) / (rewards.std())\n",
    "array([ 1.22474487,  1.22474487, -0.81649658, -0.81649658, -0.81649658])\n",
    "```\n",
    "\n",
    "Here, the correct responses receive higher advantage scores, promoting them in future updates.\n",
    "\n",
    "\n",
    "If only one response is correct:\n",
    "\n",
    "```python\n",
    ">>> rewards = np.array([1, 0, 0, 0, 0])\n",
    ">>> (rewards - rewards.mean()) / (rewards.std())\n",
    "array([ 2. , -0.5, -0.5, -0.5, -0.5])\n",
    "```\n",
    "\n",
    "This resembles the case where the question in the prompt is too hard and the model is not able to generate a correct response on average.\n",
    "However, if one of the responses is correct, it will be assigned a higher advantage score, and all incorrect responses will be assigned a negative relative score.\n",
    "\n",
    "If all responses are incorrect:\n",
    "\n",
    "```python\n",
    ">>> rewards = np.array([0, 0, 0, 0, 0])\n",
    ">>> (rewards - rewards.mean()) / (rewards.std() + 1e-6)\n",
    "array([0., 0., 0., 0., 0.])\n",
    "```\n",
    "\n",
    "Since there is no one is better than the average, the model receives no learning signal.\n",
    "\n",
    "If all responses are correct:\n",
    "\n",
    "```python\n",
    ">>> rewards = np.array([1, 1, 1, 1, 1])\n",
    ">>> (rewards - rewards.mean()) / (rewards.std() + 1e-6)\n",
    "array([0., 0., 0., 0., 0.])\n",
    "```\n",
    "\n",
    "Again, no learning signal is provided because there is nothing to improve upon.\n",
    "\n",
    "In a more mixed case:\n",
    "\n",
    "```python\n",
    ">>> rewards = np.array([1, 1, 1, 1, 0])\n",
    ">>> (rewards - rewards.mean()) / (rewards.std() + 1e-6)\n",
    "array([0.5, 0.5, 0.5, 0.5, -2.])\n",
    "```\n",
    "\n",
    "This represents an easier question for the model. Most responses are correct, but occasional incorrect ones are heavily penalized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_training_episodes(\n",
    "    samples: List[Dict[str, Any]],\n",
    "    all_generations: List[List[int]],\n",
    "    all_finish_reasons: List[str],\n",
    ") -> Tuple[Dict[str, Any], Dict[str, Any]]:\n",
    "    \"\"\"\n",
    "    Process model generations and calculate rewards for training episodes.\n",
    "\n",
    "    This function processes generated responses and calculates rewards for training episodes by:\n",
    "    1. Grouping generations by sample (GENERATIONS_PER_SAMPLE responses per input)\n",
    "    2. Computing rewards and advantages for each response\n",
    "    3. Processing response tokens\n",
    "\n",
    "    Args:\n",
    "        samples: List of input samples, each containing:\n",
    "            - input_ids: List[int], tokenized input prompt\n",
    "            - nums: List[int], numbers to use in equation\n",
    "            - target: int, target value for equation\n",
    "        all_generations: List of token ID sequences for each generated response\n",
    "        all_finish_reasons: List of finish reasons for each generation (\"stop\" or other)\n",
    "\n",
    "    Returns:\n",
    "        Tuple containing:\n",
    "        1. Dictionary with processed data for training:\n",
    "            - all_query_token_ids: List[List[int]], input token IDs repeated for each generation\n",
    "            - all_response_token_ids: List[List[int]], response token IDs with EOS tokens added\n",
    "            - all_advantages: List[List[float]], advantage values repeated for each token\n",
    "        2. Dictionary with generation statistics:\n",
    "            - response_lengths: List[int], lengths of generated responses\n",
    "            - rewards: List[float], raw reward values\n",
    "            - non_stop_rate: List[bool], whether each generation ended naturally\n",
    "            - reward_metrics/*: Various reward component metrics\n",
    "\n",
    "    Example:\n",
    "        >>> samples = [{\"input_ids\": [1,2,3], \"nums\": [1,2,3], \"target\": 6}]\n",
    "        >>> generations = [[4,5, EOS_TOKEN_ID], [6,7], [8,9, EOS_TOKEN_ID]]  # 3 generations per sample\n",
    "        >>> finish_reasons = [\"stop\", \"length\", \"stop\"]\n",
    "        >>> episodes, stats = create_training_episodes(samples, generations, finish_reasons)\n",
    "        >>> episodes\n",
    "        {\n",
    "            'all_query_token_ids': [[1,2,3], [1,2,3], [1,2,3]],\n",
    "            'all_response_token_ids': [[4,5,EOS_TOKEN_ID], [6,7], [8,9,EOS_TOKEN_ID]],\n",
    "            'all_advantages': [[0.5,0.5,0.5], [-1.0,-1.0], [0.5,0.5,0.5]]\n",
    "        }\n",
    "    \"\"\"\n",
    "    assert len(all_generations) == len(all_finish_reasons)\n",
    "    assert len(all_generations) == len(samples) * GENERATIONS_PER_SAMPLE\n",
    "\n",
    "    # Process responses and calculate rewards\n",
    "    groups = [\n",
    "        list(range(i, i + GENERATIONS_PER_SAMPLE))\n",
    "        for i in range(0, len(all_generations), GENERATIONS_PER_SAMPLE)\n",
    "    ]  # example: [[0, 1, 2], [3, 4, 5], [6, 7, 8]]\n",
    "\n",
    "    all_query_token_ids, all_responses_token_ids, all_advantages = [], [], []\n",
    "\n",
    "    stats = {\n",
    "        \"response_lengths\": [],\n",
    "        \"rewards\": [],\n",
    "        \"non_stop_rate\": [],\n",
    "    }\n",
    "\n",
    "    for sample, group_indices in zip(samples, groups):\n",
    "        finish_reasons = [all_finish_reasons[i] for i in group_indices]\n",
    "        response_token_ids = [all_generations[i] for i in group_indices]\n",
    "        responses = tokenizer.batch_decode(response_token_ids, skip_special_tokens=False)\n",
    "\n",
    "        rewards_and_metrics = [compute_reward(resp, sample) for resp in responses]\n",
    "        rewards, reward_metrics = zip(*rewards_and_metrics)\n",
    "\n",
    "        rewards = np.array(rewards) # [group_size]\n",
    "        response_advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-4)\n",
    "        \n",
    "        advantages = [\n",
    "            [resp_adv] * len(resp) \n",
    "            for resp_adv, resp in zip(response_advantages, response_token_ids)\n",
    "        ]\n",
    "\n",
    "        all_query_token_ids.extend([sample[\"input_ids\"]] * GENERATIONS_PER_SAMPLE)\n",
    "        all_responses_token_ids.extend(response_token_ids)\n",
    "        all_advantages.extend(advantages)\n",
    "\n",
    "        stats[\"rewards\"].extend(rewards)\n",
    "        stats[\"non_stop_rate\"].extend([fr != \"stop\" for fr in finish_reasons])\n",
    "        stats[\"response_lengths\"].extend([len(ids) for ids in response_token_ids])\n",
    "        for rm in reward_metrics:\n",
    "            for k, v in rm.items():\n",
    "                stats.setdefault(f\"reward_metrics/{k}\", []).append(v)\n",
    "\n",
    "    episodes = {\n",
    "        \"all_query_token_ids\": all_query_token_ids,\n",
    "        \"all_response_token_ids\": all_responses_token_ids,\n",
    "        \"all_advantages\": all_advantages,\n",
    "    }\n",
    "\n",
    "    return episodes, stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case_0 = {\n",
    "    \"sample\": {\"input_ids\": [1,2,3], \"nums\": [1,2,3], \"target\": 6},\n",
    "    \"generations\": [[4,5, 22, 33], [6,7], [8,9, 11], [10,11]],\n",
    "    \"finish_reasons\": [\"stop\", \"length\", \"stop\", \"stop\"]\n",
    "}\n",
    "\n",
    "case = case_0\n",
    "episodes, stats = create_training_episodes([case[\"sample\"]], case[\"generations\"], case[\"finish_reasons\"])\n",
    "episodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case_1 = {\n",
    "    \"sample\": {\"input_ids\": [33, 44], \"nums\": [11, 7, 8], \"target\": 26},\n",
    "    \"generations\": [[1,2], [3,4], [5,6], [7,8]],\n",
    "    \"finish_reasons\": [\"stop\", \"stop\", \"length\", \"stop\"]\n",
    "}\n",
    "case = case_1\n",
    "episodes, stats = create_training_episodes([case[\"sample\"]], case[\"generations\"], case[\"finish_reasons\"])\n",
    "episodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "case_2 = {\n",
    "    \"sample\": {\"input_ids\": [9, 8, 7, 6, 5, 4], \"nums\": [1,2,3,4], \"target\": 10},\n",
    "    \"generations\": [[9,10], [11,12], [13,14], [15,16]],\n",
    "    \"finish_reasons\": [\"length\", \"length\", \"stop\", \"stop\"]\n",
    "}\n",
    "case = case_2\n",
    "episodes, stats = create_training_episodes([case[\"sample\"]], case[\"generations\"], case[\"finish_reasons\"])\n",
    "episodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the `input_ids` of this single exmaple is repeated in all of generated episodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policy Gradient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Now that we have a batch of episodes with corresponding advantages, we can compute the **policy gradient loss** to update the model.\n",
    "\n",
    "GRPO uses the same loss formulation as PPO, but the key difference lies in how advantages are computed. To understand the implementation in `compute_pg_loss`, let’s first recall the original PPO objective:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\text{PPO}} = \\mathbb{E}\\left[\\min\\left( \n",
    "\\frac{\\pi_\\theta(y_t \\mid y_{<t}, x)}{\\pi_{\\theta_{\\text{old}}}(y_t \\mid y_{<t}, x)} A_t, \\;\n",
    "\\text{clip}\\left(\n",
    "\\frac{\\pi_\\theta(y_t \\mid y_{<t}, x)}{\\pi_{\\theta_{\\text{old}}}(y_t \\mid y_{<t}, x)}, \\;\n",
    "1 - \\epsilon, \\; 1 + \\epsilon\n",
    "\\right) A_t \\right)\\right]\n",
    "$$\n",
    "\n",
    "where:\n",
    "- $ \\pi_{\\theta} $ is the current policy,\n",
    "- $ \\pi_{\\theta_{\\text{old}}} $ is the policy from the previous iteration (the policy we sampled episodes from),\n",
    "- $ A_t $ is the advantage.\n",
    "\n",
    "This objective tries to increase or decrease the probability of tokens based on the advantage $A_t$ only when the ratio between the new and old policy probabilities stays within a small range, controlled by the clipping threshold $\\epsilon$. This clipping mechanism prevents large, destabilizing updates during training.\n",
    "\n",
    "### Fully Online Setting: Simplifying the Objective\n",
    "\n",
    "In general PPO, multiple gradient steps might be taken using the same batch of episodes. However, in our case, we apply only **one gradient step per iteration** using freshly sampled episodes. That means:\n",
    "\n",
    "- $ \\pi_{\\theta} = \\pi_{\\theta_{\\text{old}}} $\n",
    "- Consequently,  \n",
    "  $$\n",
    "  \\frac{\\pi_\\theta(y_t \\mid y_{<t}, x)}{\\pi_{\\theta_{\\text{old}}}(y_t \\mid y_{<t}, x)} = 1\n",
    "  $$\n",
    "  \n",
    "Since the ratio is exactly 1:\n",
    "- The clipping function becomes inactive.\n",
    "- The $\\min(\\cdot,\\cdot)$ operator simply returns the unclipped term.\n",
    "\n",
    "So, the objective simplifies **to**:\n",
    "\n",
    "$$\n",
    "\\mathcal{L}_{\\text{PPO}} = \\mathbb{E}\\left[ \\frac{\\pi_\\theta(y_t \\mid y_{<t}, x)}{\\pi_{\\theta_{\\text{old}}}(y_t \\mid y_{<t}, x)} A_t \\right]\n",
    "$$\n",
    "\n",
    "\n",
    "Taking the gradient of this loss with respect to $\\theta$, we get:\n",
    "\n",
    "$$\n",
    "\\vec{g}_{\\text{PPO}} = \\nabla_\\theta \\mathcal{L}_{\\text{PPO}} = 2 \\underbrace{\\mathbb{E}\\left[ \\nabla_\\theta \\log \\pi_\\theta(y_t \\mid y_{<t}, x) \\cdot A_t \\right]}_{\\text{vanilla policy gradient with advantage}}\n",
    "$$\n",
    "\n",
    "This is the **standard policy gradient** formula, where the log-probabilities are weighted by the advantage. In effect, we recover vanilla REINFORCE-style learning.\n",
    "\n",
    "> Note: The a constant multiplier (like 2) does not affect the direction of the gradient and can be safely ignored.\n",
    "\n",
    "In fact, this behavior is not unique to GRPO. In all methods such as PPO, TRPO the very first gradient step after collecting new data will always reduce to this same form. Only after the optimization step the clipping or trust region constraint start to take effect.\n",
    "\n",
    "### KL Penalty\n",
    "\n",
    "The final loss also has a **KL penalty** term to ensure the new policy doesn't drift too far from a reference policy:\n",
    "\n",
    "$$\n",
    "\\mathcal{L} = \\mathcal{L}_{\\text{PPO}} - \\beta \\cdot \\text{KL}(\\pi_\\theta \\parallel \\pi_{\\theta_{\\text{ref}}})\n",
    "$$\n",
    "\n",
    "We estimate the KL divergence using the **k3 estimator** from [this blog post by Schulman](http://joschu.net/blog/kl-approx.html):\n",
    "\n",
    "$$\n",
    "\\text{KL}(\\pi_\\theta \\parallel \\pi_{\\theta_{\\text{ref}}}) = \\mathbb{E}\\left[\\frac{\\pi_{\\theta_{\\text{ref}}}(y_t \\mid y_{<t}, x)}{\\pi_\\theta(y_t \\mid y_{<t}, x)} - \\log\\left(\\frac{\\pi_{\\theta_{\\text{ref}}}(y_t \\mid y_{<t}, x)}{\\pi_\\theta(y_t \\mid y_{<t}, x)}\\right) - 1\\right]\n",
    "$$\n",
    "\n",
    "This regularization term softly constrains the updated model to remain close to the reference.\n",
    "\n",
    "\n",
    "### GRPO vs PPO/VinePPO: Key Difference\n",
    "\n",
    "The main difference between **GRPO** and methods like **PPO/VinePPO** lies in **how the advantage is computed and applied**:\n",
    "\n",
    "- In **PPO/VinePPO**, each token/step's advantage is computed individually. This allows for fine-grained credit assignment across the sequence.\n",
    "- In **GRPO**, a **single scalar advantage** is computed for the entire response and is applied **uniformly to all tokens** in that response.\n",
    "\n",
    "This distinction is illustrated below:\n",
    "\n",
    "#### A successful response in GRPO:\n",
    "<img src=\"https://github.com/McGill-NLP/nano-aha-moment/blob/main/assets/grpo_successful.png?raw=true\" alt=\"GRPO vs PPO/VinePPO: successful response\" width=\"500\">\n",
    "\n",
    "#### A failed response in GRPO:\n",
    "<img src=\"https://github.com/McGill-NLP/nano-aha-moment/blob/main/assets/grpo_unsuccessful.png?raw=true\" alt=\"GRPO vs PPO/VinePPO: failed response\" width=\"500\">\n",
    "\n",
    "In GRPO, all tokens in a response are updated with the same magnitude. In contrast, PPO/VinePPO updates each token/step with a different advantage value:\n",
    "\n",
    "<img src=\"https://github.com/McGill-NLP/nano-aha-moment/blob/main/assets/ppo_and_vineppo.png?raw=true\" alt=\"GRPO vs PPO/VinePPO: PPO and VinePPO\" width=\"500\">\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_pg_loss(\n",
    "    policy_model: Union[DeepSpeedEngine, PreTrainedModel],\n",
    "    reference_model: Union[DeepSpeedEngine, PreTrainedModel],\n",
    "    batch: Dict[str, torch.Tensor],\n",
    "    total_response_len: int,\n",
    ") -> Tuple[torch.Tensor, Dict[str, float]]:\n",
    "    \"\"\"\n",
    "    Compute the policy gradient loss with KL penalty between policy and reference models.\n",
    "\n",
    "    This function:\n",
    "    1. Computes log probabilities for both policy and reference models\n",
    "    2. Calculates KL divergence penalty between the models\n",
    "    3. Computes policy gradient loss using advantages\n",
    "    4. Combines the losses with KL coefficient\n",
    "\n",
    "    Args:\n",
    "        policy_model: The model being trained\n",
    "        reference_model: The reference model for KL penalty calculation\n",
    "        batch: Dictionary containing:\n",
    "            - input_ids: Tensor of shape [batch_size, seq_len]\n",
    "            - attention_mask: Tensor of shape [batch_size, seq_len]\n",
    "            - labels: Tensor of shape [batch_size, seq_len] with -100 for ignored positions\n",
    "            - advantages: Tensor of shape [batch_size, seq_len]\n",
    "\n",
    "    Returns:\n",
    "        Tuple containing:\n",
    "            - loss: Combined policy gradient and KL penalty loss (scalar tensor)\n",
    "            - metrics: Dictionary with detailed loss components:\n",
    "                - policy_loss: Pure policy gradient loss\n",
    "                - kl_penalty: KL divergence penalty\n",
    "                - entropy: Policy entropy\n",
    "    \"\"\"\n",
    "    input_ids = batch[\"input_ids\"]  # [batch_size, seq_len]\n",
    "    attention_mask = batch[\"attention_mask\"]  # [batch_size, seq_len]\n",
    "    labels = batch[\"labels\"]  # [batch_size, seq_len]\n",
    "    advantages = batch[\"advantages\"]  # [batch_size, seq_len]\n",
    "\n",
    "    model_inputs = {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "        \"labels\": labels,\n",
    "    }\n",
    "\n",
    "    labels_mask = (labels[..., 1:] != -100).float()  # [batch_size, seq_len-1]\n",
    "\n",
    "    with torch.no_grad():\n",
    "        ref_logps = compute_token_log_probs(\n",
    "            reference_model, model_inputs, TEMPERATURE\n",
    "        )  # [batch_size, seq_len-1]\n",
    "\n",
    "    logps = compute_token_log_probs(policy_model, model_inputs, TEMPERATURE)  # [batch_size, seq_len-1]\n",
    "\n",
    "    kl_penalty = torch.exp(ref_logps - logps) - (ref_logps - logps) - 1  # [batch_size, seq_len-1]\n",
    "    kl_penalty = kl_penalty * labels_mask  # [batch_size, seq_len-1]\n",
    "\n",
    "    entropy = -logps.sum() / labels_mask.sum()  # scalar\n",
    "\n",
    "    policy_loss = -logps * advantages[..., 1:]  # [batch_size, seq_len-1]\n",
    "    policy_loss = policy_loss * labels_mask  # [batch_size, seq_len-1]\n",
    "\n",
    "    loss = (policy_loss + KL_COEFFICIENT * kl_penalty).sum() / total_response_len  # scalar\n",
    "\n",
    "    metrics = {\n",
    "        \"policy_loss\": policy_loss.sum().item() / total_response_len,\n",
    "        \"kl_penalty\": kl_penalty.sum().item() / total_response_len,\n",
    "        \"entropy\": entropy.item() / total_response_len,\n",
    "    }\n",
    "\n",
    "    return loss, metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before starting the RL loop, we need to set up all necessary components:\n",
    "\n",
    "- **Policy Model**: The main model that will be trained using policy gradients.\n",
    "- **Reference Model**: A frozen copy of the base model used for KL regularization.\n",
    "- **DeepSpeed**: Both models are initialized with DeepSpeed.\n",
    "- **vLLM Inference Engine**: Used for fast, batched inference during episode generation.\n",
    "- **WandB Logging**: We initialize WandB to track training metrics, hyperparameters, and checkpoints.\n",
    "\n",
    "Finally, if an existing checkpoint is detected, we automatically resume training from where it left off. \n",
    "\n",
    "Couple of remarks:\n",
    "- We move the reference to CPU and only take back to GPU during policy gradient computation. Because of the relatievely small size of the model, this moving back and forth from GPU to CPU is super fast.\n",
    "- Despite the entire training being run on a single GPU, we still use DeepSeed Zero stage 2. This is because the stage 2 comes with some optimization that avoid memory fragmentations, allowing to fully utilize GPU memory.\n",
    "- Flash Attention is required in our setup as it reduces the memory requirement of transformers from $\\mathcal{O}(n^2)$ to $\\mathcal{O}(n)$ where $n$ the sequence length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize main and reference models\n",
    "policy_model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device_map=0,\n",
    ")\n",
    "reference_model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    device_map=0,\n",
    ")\n",
    "policy_model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n",
    "\n",
    "\n",
    "# Initialize DeepSpeed engines\n",
    "policy_model, *_ = deepspeed.initialize(\n",
    "    model=policy_model,\n",
    "    config=deepspeed_config,\n",
    "    model_parameters=policy_model.parameters(),\n",
    ")\n",
    "reference_model, *_ = deepspeed.initialize(\n",
    "    model=reference_model,\n",
    "    config=ref_deepspeed_config,\n",
    ")\n",
    "\n",
    "reference_model.module.cpu()\n",
    "\n",
    "############################################\n",
    "# Initialize vLLM (Inference) engine\n",
    "############################################\n",
    "\n",
    "inference_engine = LLM(\n",
    "    model=MODEL_NAME,\n",
    "    skip_tokenizer_init=False,\n",
    "    gpu_memory_utilization=0.2,\n",
    "    enable_prefix_caching=True,\n",
    "    swap_space=1,\n",
    "    scheduling_policy=\"fcfs\",\n",
    "    dtype=torch.bfloat16,\n",
    "    max_model_len=2048,\n",
    "    enable_sleep_mode=True,\n",
    ")\n",
    "\n",
    "# Wandb for logging\n",
    "wandb.init(\n",
    "    project=\"r1-aha-moment\",\n",
    "    name=RUN_NAME,\n",
    "    config={\n",
    "        \"model_name\": MODEL_NAME,\n",
    "        \"learning_rate\": LEARNING_RATE,\n",
    "        \"num_iterations\": NUM_ITERATIONS,\n",
    "        \"episodes_per_iteration\": EPISODES_PER_ITERATION,\n",
    "        \"rollouts_per_episode\": GENERATIONS_PER_SAMPLE,\n",
    "        \"kl_coefficient\": KL_COEFFICIENT,\n",
    "        \"temperature\": TEMPERATURE,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Load checkpoint if it exists\n",
    "begin_iter = 0\n",
    "ckpt_path, ckpt_iter = find_last_checkpoint(EXP_DIR)\n",
    "if ckpt_path is not None:\n",
    "    print(f\"Resuming from checkpoint {ckpt_path} at iteration {ckpt_iter}\")\n",
    "    out = policy_model.load_checkpoint(ckpt_path / \"deepspeed\")\n",
    "    if out is None:\n",
    "        raise RuntimeError(f\"Failed to load checkpoint {ckpt_path}\")\n",
    "    begin_iter = ckpt_iter + 1\n",
    "    load_model_into_vllm(policy_model, inference_engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With everything set up, we are ready to start the main training loop. Each iteration of the loop performs the following steps:\n",
    "\n",
    "1. **Evaluation** (optional): \n",
    "Every few iterations, the model is evaluated on a test set to monitor progress.\n",
    "2. **Episode Generation**\n",
    "A batch of prompts is sampled, and multiple responses are generated for each prompt using the inference engine. Then we put the inference engine to sleep.\n",
    "3. **Reward Computation**\n",
    "Rewards and advantages for each generated episode are computed.\n",
    "4. **Policy Gradient Training**\n",
    "Using the computed advantages, we calculate the policy gradient loss and update the model parameters. The training is done using gradient accumulation to handle large batches. Note that we apply single gradient update per iteration.\n",
    "5. **Inference Engine Update**\n",
    "The inference engine is woken up and updated with the latest model weights.\n",
    "6. **Logging**\n",
    "Training and evaluation metrics are logged using WandB.\n",
    "7. **Checkpointing**\n",
    "Every 50 iterations, the model and optimizer states are saved.\n",
    "\n",
    "This loop continues until the specified number of iterations is completed.\n",
    "\n",
    "**Sleeping of vLLM**\n",
    "Before training begins, we put vLLM into sleep mode to free up its KV cache and model weights, ensuring enough GPU memory is available for policy training. After the training step is complete, vLLM is woken up, reinitializing its KV cache and preparing for the next round of sampling using the updated model parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for iteration in trange(NUM_ITERATIONS):\n",
    "    print(f\"Iteration {iteration}/{NUM_ITERATIONS}\")\n",
    "\n",
    "    metrics = {}\n",
    "\n",
    "    #########################################################\n",
    "    # Evaluation\n",
    "    #########################################################\n",
    "\n",
    "    eval_stats = None\n",
    "    if iteration % 25 == 0:\n",
    "        print(\"Evaluating on eval set...\")\n",
    "        eval_episodes, eval_stats = evaluate_on_test_set(\n",
    "            inference_engine=inference_engine,\n",
    "            test_dataset=test_dataset,\n",
    "            tokenizer=tokenizer,\n",
    "            eos_token=EOS_TOKEN,\n",
    "            eval_sampling_params=SamplingParams(\n",
    "                temperature=0.3,\n",
    "                max_tokens=1024,\n",
    "                n=1,\n",
    "                detokenize=False,\n",
    "                stop_token_ids=[EOS_TOKEN_ID],\n",
    "            ),\n",
    "            reward_func=lambda completion, sample: compute_reward(\n",
    "                completion, sample\n",
    "            ),\n",
    "        )\n",
    "        eval_episode_table = dump_episodes(\n",
    "            episodes=eval_episodes,\n",
    "            episodes_stats=eval_stats,\n",
    "            exp_dir=EXP_DIR,\n",
    "            tokenizer=tokenizer,\n",
    "            iteration=iteration,\n",
    "            is_eval=True,\n",
    "        )\n",
    "        wandb.log({\"eval/episodes\": eval_episode_table, \"iteration\": iteration})\n",
    "\n",
    "\n",
    "    #########################################################\n",
    "    # Generate Episodes\n",
    "    #########################################################\n",
    "\n",
    "    # Sample training batch\n",
    "    num_samples = EPISODES_PER_ITERATION // GENERATIONS_PER_SAMPLE\n",
    "    indices = np.random.choice(\n",
    "        len(train_dataset), size=num_samples, replace=False\n",
    "    )\n",
    "    samples = train_dataset.select(indices)\n",
    "\n",
    "    # Sample responses\n",
    "    outputs = inference_engine.generate(\n",
    "        prompt_token_ids=samples[\"input_ids\"],\n",
    "        sampling_params=SamplingParams(\n",
    "            n=GENERATIONS_PER_SAMPLE,\n",
    "            temperature=TEMPERATURE,\n",
    "            top_p=TOP_P,\n",
    "            top_k=TOP_K,\n",
    "            max_tokens=MAX_RESPONSE_TOKENS,\n",
    "            detokenize=False,\n",
    "            stop_token_ids=[EOS_TOKEN_ID],\n",
    "        )\n",
    "    )\n",
    "    all_generations = [list(g.token_ids) for out in outputs for g in out.outputs]\n",
    "    all_finish_reasons = [g.finish_reason for out in outputs for g in out.outputs]\n",
    "    inference_engine.sleep(1)\n",
    "\n",
    "    print(f\"Generated {len(all_generations)} responses\")\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    time.sleep(1)\n",
    "\n",
    "    # Process responses and calculate rewards\n",
    "    episodes, episodes_stats = create_training_episodes(\n",
    "        samples,\n",
    "        all_generations,\n",
    "        all_finish_reasons,\n",
    "    )\n",
    "    for k, v in episodes_stats.items():\n",
    "        metrics.setdefault(k, []).extend(v)\n",
    "\n",
    "    episode_table = dump_episodes(\n",
    "        episodes=episodes,\n",
    "        episodes_stats=episodes_stats,\n",
    "        exp_dir=EXP_DIR,\n",
    "        tokenizer=tokenizer,\n",
    "        iteration=iteration,\n",
    "    )\n",
    "\n",
    "    #########################################################\n",
    "    # Training\n",
    "    #########################################################\n",
    "\n",
    "    # Prepare training batch\n",
    "    model_inputs = prepare_model_inputs(\n",
    "        query_token_ids=episodes[\"all_query_token_ids\"],\n",
    "        response_token_ids=episodes[\"all_response_token_ids\"],\n",
    "        advantages=episodes[\"all_advantages\"],\n",
    "        device=\"cuda\"\n",
    "    )\n",
    "\n",
    "    # Calculate losses and update model\n",
    "    policy_model.train()\n",
    "    reference_model.module.cuda()\n",
    "    reference_model.eval()\n",
    "\n",
    "    total_response_len = (model_inputs[\"labels\"] != -100).sum().item()\n",
    "\n",
    "    for i in trange(0, EPISODES_PER_ITERATION, PER_DEVICE_BATCH_SIZE, desc=\"Gradient Accumulation\"):\n",
    "        batch = {\n",
    "            k: v[i : i + PER_DEVICE_BATCH_SIZE]\n",
    "            for k, v in model_inputs.items()\n",
    "        }\n",
    "\n",
    "        # Compute policy gradient loss\n",
    "        loss, loss_metrics = compute_pg_loss(\n",
    "            policy_model=policy_model,\n",
    "            reference_model=reference_model,\n",
    "            batch=batch,\n",
    "            total_response_len=total_response_len,\n",
    "        )\n",
    "\n",
    "        # Track metrics\n",
    "        metrics.setdefault(\"loss\", []).append(loss.item())\n",
    "        grad_norm = policy_model.get_global_grad_norm()\n",
    "        if grad_norm is not None:\n",
    "            grad_norm = grad_norm.item()\n",
    "        metrics.setdefault(\"grad_norm\", []).append(grad_norm)\n",
    "        for k, v in loss_metrics.items():\n",
    "            metrics.setdefault(k, []).append(v.item() if isinstance(v, torch.Tensor) else v)\n",
    "\n",
    "        # Backpropagation and optimization step\n",
    "        policy_model.backward(loss, scale_wrt_gas=False)\n",
    "        \n",
    "        # Free memory\n",
    "        del loss, loss_metrics\n",
    "        if policy_model.is_gradient_accumulation_boundary():\n",
    "            reference_model.module.cpu()\n",
    "\n",
    "        policy_model.step()\n",
    "\n",
    "    #########################################################\n",
    "    # Update inference engine weights\n",
    "    #########################################################\n",
    "    \n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    time.sleep(1)\n",
    "\n",
    "    inference_engine.wake_up()\n",
    "    load_model_into_vllm(policy_model, inference_engine)\n",
    "\n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    time.sleep(1)\n",
    "\n",
    "\n",
    "    #########################################################\n",
    "    # Log metrics\n",
    "    #########################################################\n",
    "\n",
    "    train_metrics = {\n",
    "        k: np.mean(v) for k, v in metrics.items() if None not in v\n",
    "    }\n",
    "    train_metrics[\"learning_rate\"] = policy_model.get_lr()[0]\n",
    "    logs = {\n",
    "        \"iteration\": iteration,\n",
    "        f\"episodes/iter_{iteration:06d}\": episode_table,\n",
    "        **{f\"train/{k}\": v for k, v in train_metrics.items()},\n",
    "    }\n",
    "    if eval_stats is not None:\n",
    "        eval_metrics = {k: np.mean(v) for k, v in eval_stats.items() if None not in v}\n",
    "        logs.update({f\"eval/{k}\": v for k, v in eval_metrics.items()})\n",
    "    wandb.log(logs)\n",
    "\n",
    "    selected_keys = [\n",
    "        \"train/kl_penalty\",\n",
    "        \"train/rewards\",\n",
    "        \"train/reward_metrics/format_reward\",\n",
    "        \"train/reward_metrics/equation_reward\",\n",
    "        \"eval/rewards\",\n",
    "        \"eval/reward_metrics/format_reward\",\n",
    "        \"eval/reward_metrics/equation_reward\",\n",
    "    ]\n",
    "    selected_metrics = {k: logs[k] for k in selected_keys if k in logs}\n",
    "    print(f\"KEY METRICS: {selected_metrics}\")\n",
    "\n",
    "    if iteration % 50 == 0 and iteration != 0:\n",
    "        policy_model.module.save_pretrained(\n",
    "            str(EXP_DIR / \"checkpoints\" / f\"ckpt_{iteration:06d}\" / \"hf_model\")\n",
    "        )\n",
    "        policy_model.save_checkpoint(\n",
    "            str(EXP_DIR / \"checkpoints\" / f\"ckpt_{iteration:06d}\" / \"deepspeed\")\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Citation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you use this codebase in your research, please cite us using:\n",
    "\n",
    "```bibtex\n",
    "@misc{Kazemnejad2025:NanoAhaMoment,\n",
    "  author       = {Amirhossein Kazemnejad and Milad Aghajohari and Alessandro Sordoni and Aaron Courville and Siva Reddy},\n",
    "  title        = {Nano Aha! Moment: Lunch Break Reproduction of DeepSeek R1-Zero from Scratch},\n",
    "  year         = {2025},\n",
    "  howpublished = {\\url{https://github.com/McGill-NLP/nano-aha-moment}},\n",
    "  note         = {GitHub repository}\n",
    "}\n",
    "```"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 fixed modules",
   "language": "python",
   "name": "python3_fixed"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
